{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['C:\\\\Users\\\\sinyer\\\\AppData\\\\Local\\\\conda\\\\conda\\\\envs\\\\gluon\\\\lib\\\\site-packages\\\\ipykernel_launcher.py']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "import threading\n",
    "import numpy as np\n",
    "from pysc2.env import sc2_env\n",
    "from pysc2.lib import actions\n",
    "from pysc2.lib import features\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.layers as layers\n",
    "\n",
    "tf.set_random_seed(1)\n",
    "config = tf.ConfigProto(allow_soft_placement=True)  # auto distribute device\n",
    "config.gpu_options.allow_growth = True  # gpu memory dependent on require\n",
    "\n",
    "import sys\n",
    "from absl import flags\n",
    "flags.DEFINE_string('f', '', 'kernel')\n",
    "flags.FLAGS(sys.argv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def preprocess_minimap(feature_minimap):  # TODO 两个preprocess函数待解析\n",
    "    minimap = np.array(feature_minimap, dtype=np.float32)\n",
    "    layers = []\n",
    "    for i in range(len(features.MINIMAP_FEATURES)):\n",
    "        if i == features.MINIMAP_FEATURES.player_id.index:\n",
    "            layers.append(minimap[i:i + 1] / features.MINIMAP_FEATURES[i].scale)\n",
    "        elif features.MINIMAP_FEATURES[i].type == features.FeatureType.SCALAR:\n",
    "            layers.append(minimap[i:i + 1] / features.MINIMAP_FEATURES[i].scale)\n",
    "        else:\n",
    "            layer = np.zeros([features.MINIMAP_FEATURES[i].scale, minimap.shape[1], minimap.shape[2]], dtype=np.float32)\n",
    "            for j in range(features.MINIMAP_FEATURES[i].scale):\n",
    "                indy, indx = (minimap[i] == j).nonzero()\n",
    "                layer[j, indy, indx] = 1\n",
    "            layers.append(layer)\n",
    "    return np.expand_dims(np.concatenate(layers, axis=0), axis=0)\n",
    "\n",
    "\n",
    "def preprocess_screen(feature_screen):\n",
    "    screen = np.array(feature_screen, dtype=np.float32)\n",
    "    layers = []\n",
    "    for i in range(len(features.SCREEN_FEATURES)):\n",
    "        if i == features.SCREEN_FEATURES.player_id.index or i == features.SCREEN_FEATURES.unit_type.index:\n",
    "            layers.append(screen[i:i + 1] / features.SCREEN_FEATURES[i].scale)\n",
    "        elif features.SCREEN_FEATURES[i].type == features.FeatureType.SCALAR:\n",
    "            layers.append(screen[i:i + 1] / features.SCREEN_FEATURES[i].scale)\n",
    "        else:\n",
    "            layer = np.zeros([features.SCREEN_FEATURES[i].scale, screen.shape[1], screen.shape[2]], dtype=np.float32)\n",
    "            for j in range(features.SCREEN_FEATURES[i].scale):\n",
    "                indy, indx = (screen[i] == j).nonzero()\n",
    "                layer[j, indy, indx] = 1\n",
    "            layers.append(layer)\n",
    "    return np.expand_dims(np.concatenate(layers, axis=0), axis=0)\n",
    "\n",
    "\n",
    "class A3C:\n",
    "    def __init__(self, sess, reuse):\n",
    "        self.action_num = len(actions.FUNCTIONS)\n",
    "        \n",
    "        # 每个agent单独的会话、单独的观测传参\n",
    "        self.sess = sess\n",
    "        self.minimap = tf.placeholder(tf.float32, [None, 17, 64, 64])\n",
    "        self.screen = tf.placeholder(tf.float32, [None, 42, 64, 64])\n",
    "        self.info = tf.placeholder(tf.float32, [None, self.action_num])\n",
    "\n",
    "        self.spatial_mask = tf.placeholder(tf.float32, [None])\n",
    "        self.spatial_choose = tf.placeholder(tf.float32, [None, 64 ** 2])\n",
    "        self.non_spatial_mask = tf.placeholder(tf.float32, [None, self.action_num])\n",
    "        self.non_spatial_choose = tf.placeholder(tf.float32, [None, self.action_num])\n",
    "        self.q_target_value = tf.placeholder(tf.float32, [None])\n",
    "\n",
    "        with tf.variable_scope('agent') and tf.device('/gpu:0'):\n",
    "            if reuse:\n",
    "                tf.get_variable_scope().reuse_variables()\n",
    "\n",
    "            # ———————————————— 特征提取网络 —————————————————— #\n",
    "            \n",
    "            mconv1 = layers.conv2d(tf.transpose(self.minimap, [0, 2, 3, 1]), 16, 5, scope='mconv1')\n",
    "            mconv2 = layers.conv2d(mconv1, 32, 3, scope='mconv2')\n",
    "            sconv1 = layers.conv2d(tf.transpose(self.screen, [0, 2, 3, 1]), 16, 5, scope='sconv1')\n",
    "            sconv2 = layers.conv2d(sconv1, 32, 3, scope='sconv2')\n",
    "            info_feature = layers.fully_connected(layers.flatten(self.info), 256, activation_fn=tf.tanh,\n",
    "                                                  scope='info_feature')\n",
    "\n",
    "            flatten_concat = tf.concat([layers.flatten(mconv2), layers.flatten(sconv2), info_feature], axis=1)\n",
    "            flatten_feature = layers.fully_connected(flatten_concat, 256, activation_fn=tf.nn.relu,\n",
    "                                                    scope='flatten_feature')\n",
    "\n",
    "            # TODO 可能加入info向量信息到conv层\n",
    "            conv_concat = tf.concat([mconv2, sconv2], axis=3)\n",
    "            conv_feature = layers.conv2d(conv_concat, 1, 1, activation_fn=None, scope='conv_feature')\n",
    "\n",
    "            # ———————————————— 动作选择输出网络 —————————————————— #\n",
    "\n",
    "            self.q_value = tf.reshape(layers.fully_connected(flatten_feature, 1, activation_fn=None,\n",
    "                                                           scope='q_value'), [-1])  # TODO 作用未知\n",
    "\n",
    "            self.spatial_action = tf.nn.softmax(layers.flatten(conv_feature))\n",
    "\n",
    "            self.non_spatial_action = layers.fully_connected(flatten_feature, self.action_num, activation_fn=tf.nn.softmax,\n",
    "                                                             scope='non_spatial_action')\n",
    "\n",
    "            # ———————————————— 策略提升网络 —————————————————— #\n",
    "            # TODO 需要这部分所使用算法的详细解析\n",
    "            advantage = tf.stop_gradient(self.q_target_value - self.q_value)\n",
    "\n",
    "            # 求动作选择的似然概率\n",
    "            spatial_prob = tf.reduce_sum(self.spatial_action * self.spatial_choose, axis=1)\n",
    "            spatial_log_prob = tf.log(tf.clip_by_value(spatial_prob, 1e-10, 1.))\n",
    "\n",
    "            non_spatial_prob = tf.reduce_sum(self.non_spatial_action * self.non_spatial_choose, axis=1)\n",
    "            valid_non_spatial_prob = tf.reduce_sum(self.non_spatial_action * self.non_spatial_mask, axis=1)\n",
    "            valid_non_spatial_prob = tf.clip_by_value(valid_non_spatial_prob, 1e-10, 1.)\n",
    "            non_spatial_prob = non_spatial_prob / valid_non_spatial_prob\n",
    "            non_spatial_log_prob = tf.log(tf.clip_by_value(non_spatial_prob, 1e-10, 1.))\n",
    "\n",
    "            action_log_prob = self.spatial_mask * spatial_log_prob + non_spatial_log_prob\n",
    "\n",
    "            # 策略损失与价值损失\n",
    "            policy_loss = - tf.reduce_mean(action_log_prob * advantage)\n",
    "            value_loss = 0.25 * tf.reduce_mean(tf.square(self.q_value * advantage))\n",
    "\n",
    "            # ——————— 加入entropy loss ——————— #\n",
    "            # entropy = - (tf.reduce_sum(spatial_prob * spatial_log_prob, axis=-1) +\n",
    "            #              tf.reduce_sum(non_spatial_prob * non_spatial_log_prob, axis=-1))\n",
    "            # entropy_loss = - 1e-3 * tf.reduce_mean(entropy)\n",
    "            # loss = policy_loss + value_loss + entropy_loss\n",
    "\n",
    "            loss = policy_loss + value_loss\n",
    "\n",
    "            # ———————————————— 训练定义 —————————————————— #\n",
    "\n",
    "            opt = tf.train.RMSPropOptimizer(5e-4, decay=0.99, epsilon=1e-10)\n",
    "            grads = opt.compute_gradients(loss)\n",
    "            cliped_grad = []\n",
    "            for grad, var in grads:\n",
    "                grad = tf.clip_by_norm(grad, 10.0)\n",
    "                cliped_grad.append([grad, var])\n",
    "            self.train_op = opt.apply_gradients(cliped_grad)\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        minimap = preprocess_minimap(state.observation['feature_minimap'])\n",
    "        screen = preprocess_screen(state.observation['feature_screen'])\n",
    "\n",
    "        info = np.zeros([1, self.action_num], dtype=np.float32)\n",
    "        info[0, state.observation['available_actions']] = 1  # TODO 未知设置\n",
    "\n",
    "        feed = {self.minimap: minimap, self.screen: screen, self.info: info}\n",
    "\n",
    "        non_spatial_action, spatial_action = self.sess.run(\n",
    "            [self.non_spatial_action, self.spatial_action],\n",
    "            feed_dict=feed)\n",
    "\n",
    "        non_spatial_action = non_spatial_action.ravel()  # TODO ravel函数作用待解释\n",
    "        spatial_action = spatial_action.ravel()\n",
    "\n",
    "        valid_actions = state.observation['available_actions']\n",
    "\n",
    "        # ——————— 加入epsilon随机值衰减 ——————— #\n",
    "        # epsilon = [0.05, 0.2]\n",
    "        # if counter >= 2000:\n",
    "        #     epsilon[0] = epsilon[0] / 1.005\n",
    "        #     epsilon[1] = epsilon[1] / 1.005\n",
    "        #\n",
    "        # if np.random.rand() < epsilon[0]:\n",
    "        #     action_id = np.random.choice(valid_actions)\n",
    "        # if np.random.rand() < epsilon[1]:\n",
    "        #     noise = np.random.randint(-4, 5)\n",
    "        #     location[0] = int(max(0, min(64 - 1, location[0] + noise)))\n",
    "        #     location[1] = int(max(0, min(64 - 1, location[1] + noise)))\n",
    "        \n",
    "        action_id = valid_actions[np.argmax(non_spatial_action[valid_actions])]\n",
    "        if np.random.rand() < 0.05:  # 随机选择动作\n",
    "            action_id = np.random.choice(valid_actions)\n",
    "\n",
    "        location = np.argmax(spatial_action)\n",
    "        location = [int(location // 64), int(location % 64)]\n",
    "        if np.random.rand() < 0.2:  # TODO 随机选择坐标?\n",
    "            noise = np.random.randint(-4, 5)\n",
    "            location[0] = int(max(0, min(64 - 1, location[0] + noise)))\n",
    "            location[1] = int(max(0, min(64 - 1, location[1] + noise)))\n",
    "\n",
    "        action_args = []\n",
    "        for arg in actions.FUNCTIONS[action_id].args:\n",
    "            if arg.name in ('screen', 'minimap', 'screen2'):\n",
    "                action_args.append([location[1], location[0]])\n",
    "            else:\n",
    "                action_args.append([0])  # TODO 未知该参数意义\n",
    "                \n",
    "        print(action_args)        \n",
    "        return action_id, action_args\n",
    "\n",
    "    def update(self, buffer, epoch):\n",
    "        last_state = buffer[-1][-1]\n",
    "\n",
    "        if last_state.last():\n",
    "            R = 0  # TODO R设置的具体意义？\n",
    "        else:\n",
    "            minimap = preprocess_minimap(last_state.observation['feature_minimap'])\n",
    "            screen = preprocess_screen(last_state.observation['feature_screen'])\n",
    "            info = np.zeros([1, self.action_num], dtype=np.float32)\n",
    "            info[0, last_state.observation['available_actions']] = 1\n",
    "\n",
    "            feed = {self.minimap: minimap,\n",
    "                    self.screen: screen,\n",
    "                    self.info: info}\n",
    "            \n",
    "            R = self.sess.run(self.q_value, feed_dict=feed)[0]\n",
    "\n",
    "        minimaps = []\n",
    "        screens = []\n",
    "        infos = []\n",
    "        rewards = []\n",
    "\n",
    "        # ————— 原计算target value的设计 ————————— #\n",
    "        # target_value = np.zeros([len(buffer)], dtype=np.float32)\n",
    "        # target_value[-1] = R\n",
    "\n",
    "        spatial_mask = np.zeros([len(buffer)], dtype=np.float32)\n",
    "        spatial_choose = np.zeros([len(buffer), 64 ** 2], dtype=np.float32)\n",
    "        non_spatial_mask = np.zeros([len(buffer), self.action_num], dtype=np.float32)\n",
    "        non_spatial_choose = np.zeros([len(buffer), self.action_num], dtype=np.float32)\n",
    "\n",
    "        for i, [state, action_id, action_args, _] in enumerate(buffer):\n",
    "            # 求解minimap、screen、info、reward\n",
    "            minimap = preprocess_minimap(state.observation['feature_minimap'])\n",
    "            screen = preprocess_screen(state.observation['feature_screen'])\n",
    "            info = np.zeros([1, self.action_num], dtype=np.float32)\n",
    "            info[0, state.observation['available_actions']] = 1\n",
    "\n",
    "            reward = state.reward\n",
    "\n",
    "            minimaps.append(minimap)\n",
    "            screens.append(screen)\n",
    "            infos.append(info)\n",
    "            rewards.append(reward)\n",
    "\n",
    "            # ——————— 原计算target value的设计 ————————— #\n",
    "            # target_value[i] = reward + 0.99 * target_value[i - 1]\n",
    "\n",
    "            # 求解mask和choose\n",
    "            valid_actions = state.observation[\"available_actions\"]\n",
    "            non_spatial_mask[i, valid_actions] = 1\n",
    "            non_spatial_choose[i, action_id] = 1\n",
    "\n",
    "            args = actions.FUNCTIONS[action_id].args  # TODO 这部分具体意义待测试\n",
    "            for arg, act_arg in zip(args, action_args):\n",
    "                if arg.name in ('screen', 'minimap', 'screen2'):\n",
    "                    ind = act_arg[1] * 64 + act_arg[0]\n",
    "                    spatial_mask[i] = 1\n",
    "                    spatial_choose[i, ind] = 1\n",
    "\n",
    "        minimaps = np.concatenate(minimaps, axis=0)\n",
    "        screens = np.concatenate(screens, axis=0)\n",
    "        infos = np.concatenate(infos, axis=0)\n",
    "\n",
    "        # 求解q_target_value\n",
    "        q_target_value = np.zeros([len(buffer)], dtype=np.float32)\n",
    "        running_add = R\n",
    "        for j in reversed(range(0, len(rewards))):\n",
    "            running_add = rewards[j] + 0.99 * running_add\n",
    "            q_target_value[j] = running_add\n",
    "\n",
    "        feed = {self.minimap: minimaps,\n",
    "                self.screen: screens,\n",
    "                self.info: infos,\n",
    "                self.q_target_value: q_target_value,\n",
    "                self.spatial_mask: spatial_mask,\n",
    "                self.spatial_choose: spatial_choose,\n",
    "                self.non_spatial_mask: non_spatial_mask,\n",
    "                self.non_spatial_choose: non_spatial_choose}\n",
    "        self.sess.run(self.train_op, feed_dict=feed)\n",
    "\n",
    "        if sum(rewards) > 2:\n",
    "            print(epoch)\n",
    "\n",
    "\n",
    "def build_env(map_name):\n",
    "    if map_name == 'Simple64':\n",
    "        players = [sc2_env.Agent(sc2_env.Race['terran']),\n",
    "                   sc2_env.Bot(sc2_env.Race['terran'], sc2_env.Difficulty['very_easy'])]\n",
    "    else:\n",
    "        players = [sc2_env.Agent(sc2_env.Race['terran'])]\n",
    "    interface = sc2_env.parse_agent_interface_format(feature_screen=64, feature_minimap=64)\n",
    "    env = sc2_env.SC2Env(map_name=map_name, players=players, step_mul=8, agent_interface_format=interface)\n",
    "    return env\n",
    "\n",
    "\n",
    "def run(agent, max_epoch, map_name):\n",
    "    env = build_env(map_name)\n",
    "    buffer = []\n",
    "    \n",
    "    for epoch in range(max_epoch):\n",
    "        state = env.reset()[0]\n",
    "        counter = 0\n",
    "        max_step = 60\n",
    "        while True:\n",
    "            counter += 1\n",
    "\n",
    "            action_id, action_args = agent.choose_action(state)\n",
    "            next_state = env.step([actions.FunctionCall(action_id, action_args)])[0]\n",
    "\n",
    "            buffer.append([state, action_id, action_args, next_state])\n",
    "\n",
    "            if epoch > 4000:\n",
    "                max_step = 240\n",
    "            if counter >= max_step or next_state.last():  # TODO buffer学习的判断具体设置（num_frames >= max_frames）\n",
    "                agent.update(buffer, epoch)  # TODO 是否加入lr衰减\n",
    "                buffer = []\n",
    "                break\n",
    "\n",
    "            state = next_state\n",
    "    \n",
    "    env.close()\n",
    "            \n",
    "\n",
    "def run_a3c(max_epoch, map_name, parallel):\n",
    "    sess = tf.Session(config=config)\n",
    "\n",
    "    agents = []\n",
    "    for i in range(parallel):\n",
    "        agent = A3C(sess, i > 0)\n",
    "        agents.append(agent)\n",
    "        \n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    threads = []\n",
    "    for i in range(parallel):\n",
    "        t = threading.Thread(target=run, args=(agents[i], max_epoch, map_name))\n",
    "        threads.append(t)\n",
    "        t.daemon = True\n",
    "        t.start()\n",
    "        time.sleep(5)\n",
    "\n",
    "    coord = tf.train.Coordinator()\n",
    "    coord.join(threads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epoch = 40000\n",
    "map_name = 'MoveToBeacon'\n",
    "# parallel = 4  # GeForce 1080Ti\n",
    "parallel = 1  # GeForce 1070\n",
    "\n",
    "run_a3c(max_epoch, map_name, parallel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
